# -*- coding: utf-8 -*-

import os
import sys
import argparse
import torch

from load_data import load_vocabs, reverse_k_v
from load_data import TrainIteratorWrapper, TestIteratorWrapper
from train_eval_test import train, eval, test, post_test
from model.lstm_crf import LSTM_CRF

if __name__ == "__main__":
    if len(sys.argv) < 7:
        print("NER Model for entity type inference. Run: \n"
              "  python3 main.py train/eval train_file eval_file test_file char_vocabs ner_type_vocabs \n")
        exit(0)

    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', action="store", type=str, default="./checkpoint/models.epochs", help="select a checkpoint to restore.")
    parser.add_argument('--save-path', action="store", type=str, default="./checkpoint/models.epoch{}")
    parser.add_argument('--use-cuda', action='store', type=bool, default=True, help="whether use cuda.")
    parser.add_argument('--out-ner', action='store', type=str, default=None, help="test mode, output the predicted ners.")
    parser.add_argument('--out-text', action='store', type=str, default=None, help='test mode, output the predicted text.')
    parser.add_argument('--out-validate', action='store', type=str, default=None, help='post_test mode, output the validate.')
    args = parser.parse_args(sys.argv[7:])

    action = sys.argv[1]
    UNK, PAD = "<UNK>", "<PAD>"

    train_file = sys.argv[2]
    eval_file = sys.argv[3]
    test_file = sys.argv[4]
    char_vocabs_file = sys.argv[5]
    ner_type_vocabs_file = sys.argv[6]

    char_vocabs = load_vocabs(char_vocabs_file)
    ner_type_vocabs = load_vocabs(ner_type_vocabs_file)
    idx2char = reverse_k_v(char_vocabs)
    idx2tag = reverse_k_v(ner_type_vocabs)

    unk_idx = char_vocabs.get(UNK)
    pad_idx = char_vocabs.get(PAD)

    epochs = 20
    batch_size = 128  # 64
    max_len = 300    # 100
    emb_dim = 300    # 300
    hidden_size = 512  # 256
    use_cuda = args.use_cuda
    device = torch.device("cuda:0") if use_cuda else torch.device("cpu:0")

    train_iter = TrainIteratorWrapper(train_file, batch_size, max_len, char_vocabs, ner_type_vocabs, unk_idx, pad_idx,
                                 device)
    eval_iter = TrainIteratorWrapper(eval_file, batch_size, max_len, char_vocabs, ner_type_vocabs, unk_idx, pad_idx,
                                device)
    test_iter = TestIteratorWrapper(test_file, batch_size, max_len, char_vocabs, unk_idx, pad_idx, device)

    model = LSTM_CRF(len(char_vocabs), len(ner_type_vocabs), emb_dim, max_len, hidden_size, pad_idx=pad_idx,
                     use_cuda=use_cuda)

    if os.path.exists(args.checkpoint):
        model.load_state_dict(torch.load(args.checkpoint))
        print("restore models from %s ." % args.checkpoint)

    if action in "train":
        train(model, train_iter, eval_iter, 1e-3, epochs, idx2char, idx2tag, save_path=args.save_path)
    elif action in "eval":
        eval(model, eval_iter, idx2char, idx2tag)
    elif action in "test":
        if not args.out_ner or not args.out_text:
            raise Exception("out-ner and out-text must be specified in test mode!")
        test(model, test_iter, idx2char, idx2tag, args.out_ner, args.out_text)
        if args.out_validate is not None:
            post_test(args.out_ner, './data/ccks_7_1_competition_data/验证集/test.txt', args.out_validate)
    elif action in 'post_test':
        if not args.out_ner or not args.out_validate:
            raise Exception("out-ner and out-validate must be specified in post_test mode!")
        post_test(args.out_ner, './data/ccks_7_1_competition_data/验证集/test.txt', args.out_validate)
    else:
        raise Exception("Unsupported action: {}".format(action))

""" Run:
python3 main.py train ./data/train.txt ./data/eval.txt ./data/test.txt ./data/char_vocabs.txt ./data/ner_type_vocabs.txt
"""
